package org.github.manx98.jlibhac.tools.fssystem.ncautils;

import cn.hutool.core.util.HexUtil;
import cn.hutool.crypto.SecureUtil;
import org.github.manx98.jlibhac.common.Utilities;
import org.github.manx98.jlibhac.common.keys.KeySet;
import org.github.manx98.jlibhac.common.math.IByteUtil;
import org.github.manx98.jlibhac.common.math.Uint64;
import org.github.manx98.jlibhac.fs.IStorage;
import org.github.manx98.jlibhac.fssystem.DistributionType;
import org.github.manx98.jlibhac.fssystem.NcaContentType;
import org.github.manx98.jlibhac.tools.fssystem.Aes128XtsTransform;

import javax.crypto.ShortBufferException;
import java.io.IOException;
import java.math.BigInteger;
import java.nio.ByteOrder;

/**
 * @author Manx98
 * @version 1.0
 * @since 2022/10/3 15:21
 */
public class NcaHeader {
    private static final int HeaderSize = 0xC00;
    public static final int headerSectorSize = 0x200;
    public static final int HeaderSectorSize = 0x200;
    private final byte[] header;
    public final NcaVersion formatVersion;
    public final boolean isEncrypted;
    public NcaHeaderStruct headerStruct = new NcaHeaderStruct();

    public static final byte[] nca0FixedBodyKeySha256Hash = new byte[]
            {
                    (byte) 0x9A, (byte) 0xBB, (byte) 0xD2, 0x11, (byte) 0x86, 0x00, 0x21, (byte) 0x9D, 0x7A, (byte) 0xDC, 0x5B, 0x43, (byte) 0x95, (byte) 0xF8, 0x4E, (byte) 0xFD,
                    (byte) 0xFF, 0x6B, 0x25, (byte) 0xEF, (byte) 0x9F, (byte) 0x96, (byte) 0x85, 0x28, 0x18, (byte) 0x9E, 0x76, (byte) 0xB0, (byte) 0x92, (byte) 0xF0, 0x6A, (byte) 0xCB
            };
    public static final long BlockSize = 0x200;
    public Uint64 upperCounter;
    public static final int SectionCount = 4;

    public NcaHeader(KeySet keySet, IStorage headerStorage) throws IOException {
        DecryptHeaderInfo decryptHeaderInfo = decryptHeader(keySet, headerStorage);
        this.header = decryptHeaderInfo.header;
        this.isEncrypted = decryptHeaderInfo.encrypted;
        formatVersion = detectNcaVersion(header);
        IByteUtil.mapToStruct(this.headerStruct, header, 0);
    }

    private static NcaVersion detectNcaVersion(byte[] header) {
        int version = header[0x203] - '0';

        if (version == 3) {
            return NcaVersion.Nca3;
        }
        if (version == 2) {
            return NcaVersion.Nca2;
        }
        if (version != 0) {
            return NcaVersion.Unknown;
        }

        // There are multiple versions of NCA0 that each encrypt the key area differently.
        // Examine the key area to detect which version this NCA is.

        // The end of the key area will only be non-zero if it's RSA-OAEP encrypted
        if (Utilities.spansEqual(header, NcaHeaderStruct.KeyAreaOffset + 0x80, new byte[16], 0, 0x10)) {
            return NcaVersion.Nca0RsaOaep;
        }

        // Key areas using fixed, unencrypted keys always use the same keys.
        // Check for these keys by comparing the key area with the known hash of the fixed body keys.
        if (Utilities.spansEqual(header, NcaHeaderStruct.KeyAreaOffset, nca0FixedBodyKeySha256Hash, 0, 0x20)) {
            return NcaVersion.Nca0FixedKey;
        }
        // Otherwise the key area is encrypted the same as modern NCAs.
        return NcaVersion.Nca0;
    }

    public boolean hasRightsId() {
        return !Utilities.isZeros(headerStruct.RightsId);
    }

    public Uint64 getCounter() {
        return upperCounter;
    }

    public byte[] getEncryptedKey(int index) {
        if (index < 0 || index >= SectionCount) {
            throw new IllegalArgumentException("Key index must be between 0 and 3. Actual: " + index);
        }

        int offset = NcaHeaderStruct.KeyAreaOffset + 16 * index;
        byte[] bytes = new byte[16];
        System.arraycopy(header, offset, bytes, 0, 16);
        return bytes;
    }

    public NcaContentType getContentType() {
        return headerStruct.ContentType;
    }

    public NcaFsHeader getFsHeader(int index) {
        byte[] expectedHash = getFsHeaderHash(index);
        int offset = NcaHeaderStruct.FsHeadersOffset + NcaHeaderStruct.FsHeaderSize * index;
        byte[] headerData = IByteUtil.slice(header, offset, NcaHeaderStruct.FsHeaderSize);
        byte[] actualHash = SecureUtil.sha256().digest(headerData);
        if (!Utilities.spansEqual(expectedHash, 0, actualHash, 0, actualHash.length)) {
            throw new IllegalArgumentException("FS header hash is invalid.");
        }
        return new NcaFsHeader(headerData);
    }

    public Uint64 getTitleId() {
        return headerStruct.TitleId;
    }

    public byte getKeyGeneration() {
        return (byte) Math.max(this.headerStruct.KeyGeneration1, this.headerStruct.KeyGeneration2);
    }

    public byte[] getRightsId() {
        return headerStruct.RightsId;
    }

    public byte[] GetKeyArea() {
        return IByteUtil.slice(header, NcaHeaderStruct.KeyAreaOffset, NcaHeaderStruct.KeyAreaSize);
    }

    public int getKeyAreaKeyIndex() {
        return headerStruct.KeyAreaKeyIndex;
    }

    public static class DecryptHeaderInfo {
        public final byte[] header;
        public final boolean encrypted;

        public DecryptHeaderInfo(byte[] header, boolean encrypted) {
            this.header = header;
            this.encrypted = encrypted;
        }
    }

    public static DecryptHeaderInfo decryptHeader(KeySet keySet, IStorage headerStorage) throws IOException {
        byte[] header = headerStorage.readBytes(HeaderSize);
        if (checkIfDecrypted(header)) {
            int decVersion = header[0x203] - '0';

            if (decVersion != 0 && decVersion != 2 && decVersion != 3) {
                throw new IllegalArgumentException("NCA version " + decVersion + " is not supported.");
            }
            return new DecryptHeaderInfo(header, false);
        }
        byte[] key1 = keySet.HeaderKey.subKeys(0);
        byte[] key2 = keySet.HeaderKey.subKeys(1);
        Aes128XtsTransform transform = new Aes128XtsTransform(key1, key2, true);
        transform.transformBlock(header, 0, headerSectorSize, 0);
        transform.transformBlock(header, headerSectorSize, headerSectorSize, 1);
        if (header[0x200] == 'N' && header[0x201] == 'C' && header[0x202] == 'A') {
            int version = header[0x203] - '0';
            if (version == 3) {
                for (int sector = 2; sector < HeaderSize / headerSectorSize; sector++) {
                    transform.transformBlock(header, sector * headerSectorSize, headerSectorSize, sector);
                }
            } else if (version == 2) {
                for (int i = 0x400; i < HeaderSize; i += headerSectorSize) {
                    transform.transformBlock(header, i, headerSectorSize, 0);
                }
            } else if (version != 0) {
                throw new IllegalArgumentException("NCA version " + version + " is not supported.");
            }
            return new DecryptHeaderInfo(header, true);
        }
        throw new IllegalArgumentException("Invalid header: " + new String(header, 0, 5));
    }

    private static boolean checkIfDecrypted(byte[] header) {
        // Check the magic value
        if (header[0x200] != 'N' || header[0x201] != 'C' || header[0x202] != 'A') {
            return false;
        }

        // Check the version in the magic value
        if (header[0x203] < '0' || header[0x203] > '9') {
            return false;
        }

        // Is the distribution type valid?
        if (header[0x204] > DistributionType.GameCard.ordinal()) {
            return false;
        }

        // Is the content type valid?
        return header[0x205] <= NcaContentType.PublicData.ordinal();
    }

    private byte[] getFsHeaderHash(int index) {
        int offset = NcaHeaderStruct.FsHeaderHashOffset + NcaHeaderStruct.FsHeaderHashSize * index;
        return IByteUtil.slice(header, offset, NcaHeaderStruct.FsHeaderHashSize);
    }

    public boolean isSectionEnabled(int index) {
        NcaSectionEntryStruct info = getSectionEntry(index);
        int sectStart = info.startBlock;
        int sectSize = info.endBlock - sectStart;
        return sectStart != 0 || sectSize != 0;
    }

    private NcaSectionEntryStruct getSectionEntry(int index) {
        NcaSectionEntryStruct ncaSectionEntryStruct = new NcaSectionEntryStruct();
        int offset = NcaHeaderStruct.SectionEntriesOffset + NcaSectionEntryStruct.SectionEntrySize * index;
        IByteUtil.mapToStruct(ncaSectionEntryStruct, header, offset);
        return ncaSectionEntryStruct;
    }

    public boolean isNca0() {
        return formatVersion.ordinal() >= NcaVersion.Nca0.ordinal();
    }

    public long getSectionStartOffset(int index) {
        return blockToOffset(getSectionEntry(index).startBlock);
    }

    private long blockToOffset(int blockIndex) {
        return (long) blockIndex * BlockSize;
    }

    public long getSectionSize(int index) {
        NcaSectionEntryStruct info = getSectionEntry(index);
        return blockToOffset(info.endBlock - info.startBlock);
    }
}
